﻿using System.Collections.Generic;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Input;
using System.Windows.Media;

namespace shift_多选文件
{
    public static class TreeViewMultiSelect
    {
        private static readonly Dictionary<TreeView, SelectionState> states = new();

        public static bool GetIsEnabled(DependencyObject obj)
        {
            return (bool)obj.GetValue(IsEnabledProperty);
        }

        public static void SetIsEnabled(DependencyObject obj, bool value)
        {
            obj.SetValue(IsEnabledProperty, value);
        }

        public static readonly DependencyProperty IsEnabledProperty =
            DependencyProperty.RegisterAttached("IsEnabled", typeof(bool), typeof(TreeViewMultiSelect),
                new PropertyMetadata(false, OnIsEnabledChanged));

        private static void OnIsEnabledChanged(DependencyObject d, DependencyPropertyChangedEventArgs e)
        {
            if (d is TreeView tree)
            {
                if ((bool)e.NewValue)
                {
                    var state = new SelectionState(tree);
                    states[tree] = state;

                    // 监听容器生成事件
                    tree.ItemContainerGenerator.StatusChanged += (s, ev) =>
                    {
                        if (tree.ItemContainerGenerator.Status == System.Windows.Controls.Primitives.GeneratorStatus.ContainersGenerated)
                        {
                            foreach (TreeViewItem item in tree.Items)
                            {
                                AttachItem(item, state);
                            }
                        }
                    };
                }
                else
                {
                    states.Remove(tree);
                }
            }
        }

        private static void AttachItem(TreeViewItem item, SelectionState state)
        {
            item.Selected += (s, e) =>
            {
                e.Handled = true;
                state.HandleSelection(item);
            };
        }

        // 内部状态类
        private class SelectionState
        {
            private readonly TreeView tree;
            private readonly List<TreeViewItem> selectedItems = new();
            private TreeViewItem anchorItem;

            public SelectionState(TreeView tree)
            {
                this.tree = tree;
            }

            public void HandleSelection(TreeViewItem item)
            {
                bool ctrl = Keyboard.IsKeyDown(Key.LeftCtrl) || Keyboard.IsKeyDown(Key.RightCtrl);
                bool shift = Keyboard.IsKeyDown(Key.LeftShift) || Keyboard.IsKeyDown(Key.RightShift);

                if (shift && anchorItem != null)
                {
                    SelectRange(anchorItem, item);
                }
                else if (ctrl)
                {
                    if (selectedItems.Contains(item))
                        DeselectItem(item);
                    else
                        SelectItem(item);
                }
                else
                {
                    foreach (var sel in new List<TreeViewItem>(selectedItems))
                        DeselectItem(sel);

                    SelectItem(item);
                    anchorItem = item;
                }
            }

            private void SelectRange(TreeViewItem start, TreeViewItem end)
            {
                var items = new List<TreeViewItem>();
                foreach (TreeViewItem t in tree.Items)
                    items.Add(t);

                int i1 = items.IndexOf(start);
                int i2 = items.IndexOf(end);
                if (i1 < 0 || i2 < 0) return;

                int from = System.Math.Min(i1, i2);
                int to = System.Math.Max(i1, i2);

                foreach (var sel in new List<TreeViewItem>(selectedItems))
                    DeselectItem(sel);

                for (int i = from; i <= to; i++)
                    SelectItem(items[i]);
            }

            private void SelectItem(TreeViewItem item)
            {
                if (!selectedItems.Contains(item))
                    selectedItems.Add(item);
                item.Background = new SolidColorBrush(Color.FromRgb(100, 149, 237)); // 高亮
            }

            private void DeselectItem(TreeViewItem item)
            {
                selectedItems.Remove(item);
                item.Background = Brushes.Transparent;
            }
        }
    }
}
